#!/bin/bash

# 主程序脚本：串联 reindex、截包和 DR/INS 精度对比
# 支持批量处理多个bag压缩包
# 使用方法:
#   1. 零参数（使用默认路径）
#   2. 1个参数：<bag文件夹路径>
#   3. 2个参数：<bag文件夹路径> <ros2_msg文件夹路径>
#   4. 3个参数：<bag文件夹路径> <ros2_msg文件夹路径> <输出目录路径>

# -------------------------- 配置参数 --------------------------
# 定义默认路径
DEFAULT_BAG_DIR="../bag/"
# 默认ros2_msg文件夹路径（与脚本上级目录同级）
DEFAULT_ROS2_MSG_DIR="../ros2_msg"
DEFAULT_OUTPUT_DIR="../truncated"

# 定义相关脚本路径（当前目录下）
REINDEX_SCRIPT="./ros2_bag_reindex.sh"
TRUNCATE_SCRIPT="./bag_truncation.py"
EVAL_SCRIPT="./run_evaluation.py"  # 同级目录下的评测脚本
# --------------------------------------------------------------

# 检查必要工具是否存在
check_dependencies() {
    # 检查tar工具（用于解压压缩包）
    if ! command -v tar &> /dev/null; then
        echo "错误: 未找到 tar 工具，请安装tar后重试"
        exit 1
    fi

    # 检查Python环境
    if ! command -v python3 &> /dev/null; then
        echo "错误: 未找到 python3 环境"
        exit 1
    fi
}

# 检查必要脚本是否存在
check_scripts() {
    # 检查 reindex 脚本
    if [ ! -f "$REINDEX_SCRIPT" ]; then
        echo "错误: 未找到 reindex 脚本 '$REINDEX_SCRIPT'"
        exit 1
    fi
    if [ ! -x "$REINDEX_SCRIPT" ]; then
        echo "错误: reindex 脚本 '$REINDEX_SCRIPT' 没有执行权限"
        echo "请先执行: chmod +x $REINDEX_SCRIPT"
        exit 1
    fi

    # 检查截包脚本
    if [ ! -f "$TRUNCATE_SCRIPT" ]; then
        echo "错误: 未找到截包脚本 '$TRUNCATE_SCRIPT'"
        exit 1
    fi

    # 检查评测脚本
    if [ ! -f "$EVAL_SCRIPT" ]; then
        echo "错误: 未找到评测脚本 '$EVAL_SCRIPT'（需与main.sh同级）"
        exit 1
    fi

    # 检查 setup 文件
    if [ ! -f "$SETUP_PATH" ]; then
        echo "错误: 未找到 setup 脚本 '$SETUP_PATH'"
        exit 1
    fi
}

# 解压单个tar包到指定目录
extract_single_tar() {
    local tar_file="$1"
    local target_dir="$2"
    
    # 确保目标目录存在
    mkdir -p "$target_dir"
    
    # 检查文件是否有效
    if ! tar -tf "$tar_file" &> /dev/null; then
        echo "警告: $tar_file 不是有效的压缩文件，跳过"
        return 1
    fi
    
    # 解压到目标目录
    echo "开始解压: $tar_file 到 $target_dir"
    if [[ "$tar_file" == *.tar.gz ]]; then
        tar -xzf "$tar_file" -C "$target_dir"
    elif [[ "$tar_file" == *.tar.xz ]]; then
        tar -xJf "$tar_file" -C "$target_dir"
    else
        echo "警告: 不支持的压缩格式，跳过 $tar_file"
        return 1
    fi

    if [ $? -eq 0 ]; then
        echo "解压成功: $tar_file"
        return 0
    else
        echo "警告: 解压 $tar_file 失败"
        return 1
    fi
}

# 处理单个bag压缩包
process_single_bag() {
    local tar_file="$1"
    local base_name="$2"
    local temp_bag_dir="$3"
    local output_dir="$4"
    
    echo -e "\n======================================================"
    echo "开始处理压缩包: $tar_file"
    echo "临时工作目录: $temp_bag_dir"
    echo "输出目录: $output_dir"
    echo "======================================================"
    
    # 解压当前压缩包到临时目录
    if ! extract_single_tar "$tar_file" "$temp_bag_dir"; then
        echo "处理 $tar_file 失败，跳过"
        return 1
    fi
    
    # 创建当前包的输出目录
    mkdir -p "$output_dir"
    mkdir -p "$output_dir/evaluation_results"
    
    # 执行 reindex 操作
    run_reindex "$temp_bag_dir"
    
    # 检查 reindex 是否成功
    if [ $? -ne 0 ]; then
        echo "===== $base_name reindex 执行失败，跳过 ====="
        return 1
    fi
    
    # 执行截包操作
    run_truncate "$temp_bag_dir" "$output_dir"
    
    # 检查截包是否成功
    if [ $? -ne 0 ]; then
        echo "===== $base_name 截包执行失败，跳过 ====="
        return 1
    fi
    
    # 执行 DR 与 INS 精度对比
    run_evaluation "$output_dir"
    
    # 检查评测是否成功
    if [ $? -ne 0 ]; then
        echo "===== $base_name 精度对比执行失败，跳过 ====="
        return 1
    fi
    
    echo -e "\n===== $base_name 处理完成 ====="
    echo "结果路径: $output_dir/evaluation_results"
    return 0
}

# 执行 reindex 操作
run_reindex() {
    local bag_dir="$1"
    echo "===== 开始第一步：执行 ros2 bag reindex ====="
    echo "使用的 bag 路径: $bag_dir"
    echo "使用的 setup 路径: $SETUP_PATH"
    $REINDEX_SCRIPT "$bag_dir" "$SETUP_PATH"
    return $?
}

# 执行截包操作
run_truncate() {
    local bag_dir="$1"
    local output_dir="$2"
    echo -e "\n===== 开始第二步：执行 bag 截包 ====="
    echo "输入 bag 路径: $bag_dir"
    echo "输出路径: $output_dir"
    echo "使用的 setup 路径: $SETUP_PATH"

    # 查找 reindex 后的 db3 文件
    db3_files=($(find "$bag_dir" -type f -name "*.db3"))
    if [ ${#db3_files[@]} -eq 0 ]; then
        echo "错误: 在 $bag_dir 及其子目录中未找到 db3 文件"
        return 1
    fi
    
    # 使用找到的第一个 db3 文件作为输入
    input_bag="${db3_files[0]}"
    echo "使用的 bag 文件: $input_bag"

    # 执行 Python 截包脚本
    bash -c "source $SETUP_PATH && python3 $TRUNCATE_SCRIPT \
        -i "$input_bag" \
        -o "$output_dir" \
        -m "$ROS_INSTALL_PATH" \
        ${USE_SOURCE_MODE:+--use-source-mode}"
        
    return $?
}

# 执行 DR 与 INS 精度对比
run_evaluation() {
    local output_dir="$1"
    echo -e "\n===== 开始第三步：DR 与 INS 精度对比 ====="
    echo "评测输入路径（截包输出）: $output_dir"
    echo "ROS 安装路径: $ROS_INSTALL_PATH"
    echo "使用评测脚本: $EVAL_SCRIPT"

    # 加载 ROS 环境并执行评测脚本
    source "$SETUP_PATH" && \
    python3 "$EVAL_SCRIPT" \
        --msg-path "$ROS_INSTALL_PATH" \
        --dir "$output_dir" \
        --output-path "$output_dir/evaluation_results" \
        --name "replan" \
        --model "eval" \
        --pic true  \
        ${USE_SOURCE_MODE:+--use-source-mode}

    return $?
}

# 主流程
main() {
    # 检查基础依赖
    check_dependencies

    # 根据输入参数确定实际路径
    case $# in
        0)
            # 零参数：使用默认路径
            BAG_DIR="$DEFAULT_BAG_DIR"
            ROS2_MSG_DIR="$DEFAULT_ROS2_MSG_DIR"
            OUTPUT_DIR="$DEFAULT_OUTPUT_DIR"
            ;;
        1)
            # 1个参数：自定义 bag 路径，使用默认ros2_msg和输出路径
            BAG_DIR="$1"
            ROS2_MSG_DIR="$DEFAULT_ROS2_MSG_DIR"
            OUTPUT_DIR="$DEFAULT_OUTPUT_DIR"
            ;;
        2)
            # 2个参数：自定义 bag 和 ros2_msg 路径，使用默认输出路径
            BAG_DIR="$1"
            ROS2_MSG_DIR="$2"
            OUTPUT_DIR="$DEFAULT_OUTPUT_DIR"
            ;;
        3)
            # 3个参数：完全自定义路径
            BAG_DIR="$1"
            ROS2_MSG_DIR="$2"
            OUTPUT_DIR="$3"
            ;;
        *)
            # 参数过多
            echo "错误: 参数数量不正确"
            echo "使用方法:"
            echo "  $0 [bag文件夹路径] [ros2_msg文件夹路径] [输出目录路径]"
            exit 1
            ;;
    esac

    # 解压ros2_msg目录中的install压缩包
    if [ ! -d "${ROS2_MSG_DIR}/install" ]; then
        echo -e "\n===== 检查并解压install压缩包 ====="
        extract_single_tar "$ROS2_MSG_DIR"/*.tar* "$ROS2_MSG_DIR"
    else
        echo "检测到 ${ROS2_MSG_DIR}/install 已存在，跳过解压 ros2_msg"
    fi

    # 查找解压后的install文件夹
    INSTALL_DIR=$(find "$ROS2_MSG_DIR" -maxdepth 1 -type d -name "install" | head -n 1)
    if [ -z "$INSTALL_DIR" ]; then
        echo "错误: 在 $ROS2_MSG_DIR 中未找到解压后的install文件夹"
        exit 1
    fi

    # 设置setup.bash和ROS安装路径
    SETUP_PATH="$INSTALL_DIR/setup.bash"
    ROS_INSTALL_PATH="$INSTALL_DIR"
    echo "ROS 安装路径: $ROS_INSTALL_PATH"
    echo "setup.bash 路径: $SETUP_PATH"

    # 检查脚本依赖
    check_scripts

    # 查找所有bag压缩包
    echo -e "\n===== 查找所有bag压缩包 ====="
    tar_files=$(find "$BAG_DIR" -maxdepth 1 -type f \( -name "*.tar.gz" -o -name "*.tar.xz" \))
    if [ -z "$tar_files" ]; then
        echo "错误: 在 $BAG_DIR 中未发现 .tar.gz 或 .tar.xz 压缩包"
        exit 1
    fi

    # 创建总临时目录
    TEMP_ROOT=$(mktemp -d -t bag_processing.XXXXXX)
    echo "创建临时根目录: $TEMP_ROOT"

    # 统计压缩包数量
    total_count=$(echo "$tar_files" | wc -l)
    current_count=0
    success_count=0
    fail_count=0

    echo -e "\n===== 共发现 $total_count 个压缩包，开始批量处理 ====="

    # 循环处理每个压缩包
    for tar_file in $tar_files; do
        ((current_count++))
        # 获取文件名（不含路径和扩展名）
        base_name=$(basename "$tar_file" | sed -E 's/\.(tar\.gz|tar\.xz)$//')
        echo -e "\n===== 处理进度: $current_count/$total_count ====="
        
        # 为每个压缩包创建独立的临时目录和输出目录
        temp_bag_dir="$TEMP_ROOT/$base_name"
        output_dir="$OUTPUT_DIR/$base_name"
        
        # 处理当前压缩包
        if process_single_bag "$tar_file" "$base_name" "$temp_bag_dir" "$output_dir"; then
            ((success_count++))
        else
            ((fail_count++))
        fi
    done

    # 清理临时目录
    echo -e "\n===== 清理临时文件 ====="
    rm -rf "$TEMP_ROOT"

    # 输出处理结果统计
    echo -e "\n===== 所有处理完成 ====="
    echo "总处理压缩包数: $total_count"
    echo "成功: $success_count"
    echo "失败: $fail_count"
    echo "所有结果路径: $OUTPUT_DIR"
    exit 0
}

# 启动主流程
main "$@"
